import json

import pandas as pd
from pandas.core.frame import DataFrame
from tools.tool import tool
from tools.dbTool import dbTool
from tools.cacheTool import cache


class DataTool:
    # daily = 'daily_real'
    daily = "ts_daily_hfq"
    resultType = "dataFrame"

    def __init__(self):
        pass

    def getCalcTradeCode(self, startDate, endDate):
        """
        获取上市公司股票代码用于计算
        :rtype:DataFrame
        """
        key = "calcTradeCode.json"
        res = cache.get(key)
        if not res:
            sql = f"select distinct ts_code from {self.daily} where trade_date>={startDate} and trade_date<={endDate} " \
                  f"and ts_code not like '%BJ' and ts_code not like '300%'"
            df = dbTool.getDF(sql)
            res = df.ts_code
            res = list(res)
            resStr = json.dumps(res)
            cache.set(key, resStr)
        # 获取其中1/4
        i = 0
        res2 = []
        for val in res:
            i += 1
            if i % 4 == 0:
                res2.append(val)
        # return res
        return res2

    def getCodes(self, name='小市值'):
        fileName = f'file/codes/{name}.json'
        content = tool.readFile(fileName)
        res = json.loads(content)
        res = [i.strip().replace('XSHE', 'SZ').replace('XSHG', 'SH') for i in res]
        return res

    def getCompanyBasic(self, ts_code, start_date, end_date):
        """
        获取指定公司多个日期的基本面指标, 参数示例:'000001.sz', '20200101', '20200821'
        :param ts_code: 股票代码
        :param start_date: 开始日期(YYYYMMDD)
        :param end_date: 结束日期(YYYYMMDD)
        :rtype: DataFrame
        """
        tsp = self.getTsp()
        return tsp.daily_basic(
            ts_code=ts_code, start_date=start_date, end_date=end_date
        )

    def getTradeDate(self, startDate="20100101", table="ts_daily_hfq"):
        """
        通过股票历史记录获取历史交易日期
        """
        sql = "SELECT trade_date FROM {0} where trade_date>{1} group by trade_date order by trade_date desc limit 2000".format(
            table, startDate
        )
        res = dbTool.getDF(sql)
        return res.trade_date[::-1]

    def getAllTradeCode(self):
        sql = "SELECT distinct ts_code FROM {0} order by ts_code".format(self.daily)
        df = dbTool.getDF(sql)
        res = df.ts_code
        return res

    def getIndex(self):
        sql = "select * from index_basic"
        df = dbTool.getDF(sql)
        return df

    def getKdata(self, code, startDate=None, endDate=None):
        """
        :param code: 交易代码
        :param startDate: 开始日期,默认一年前
        :param endDate:	结束日期,默认当天
        :return: 交易数据
        """
        startDate_ = '20180101'
        endDate_ = '20200101'
        if endDate is None:
            endDate = endDate_
            # endDate = datetime.datetime.now().strftime("%Y%m%d")
        if startDate is None:
            startDate = startDate_
            # startDate = (datetime.datetime.now() + datetime.timedelta(days=-365)).strftime("%Y%m%d")
        sql = f"select * from ts_daily_hfq where ts_code='{code}' and trade_date>={startDate} and trade_date<={endDate} order by trade_date asc"
        df = dbTool.getDF(sql)
        return df

    def getSignals(self, code, param):
        data = dbTool.getList(
            "select id,param,detail from my_signals where code='{0}' and param='{1}'".format(
                code, param
            )
        )
        if not data:
            return {}
        data2 = {}
        for item in data:
            return item

    def result(self, sql):
        conn = dbTool.getDbConn()
        if self.resultType == "dataFrame":
            res = pd.read_sql(sql, conn)
        else:
            res = dbTool.getDF(sql)

        return res


dataTool = DataTool()
